# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     test
   Description :   
   Author :       lth
   date：          2022/8/3
-------------------------------------------------
   Change Activity:
                   2022/8/3 18:22: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import numpy as np
import torch
from PIL import Image
from torchvision import transforms

from config import GetConfig
from model import BiSeNet

predict_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])


cls=['background','skin', 'l_brow',
     'r_brow', 'l_eye',
     'r_eye', 'eye_g', 'l_ear',
     'r_ear', 'ear_r',
     'nose', 'mouth', 'u_lip',
     'l_lip', 'neck', 'neck_l',
     'cloth', 'hair', 'hat']


part_colors = [[0, 0, 0], [255, 85, 0], [255, 170, 0],
               [255, 0, 85], [255, 0, 170],
               [0, 255, 0], [85, 255, 0], [170, 255, 0],
               [0, 255, 85], [0, 255, 170],
               [0, 0, 255], [85, 85, 255], [170, 0, 255],
               [0, 115, 155], [0, 170, 255],
               [255, 255, 0], [255, 255, 85], [255, 255, 170],
               [44, 0, 255], [255, 85, 255], [255, 170, 255],
               [0, 255, 255], [85, 255, 255], [170, 255, 255]]


class Inference:
    def __init__(self):
        self.args = GetConfig()
        print(f"-----------{self.args.project_name}-------------")
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if use_cuda else "cpu")
        self.model = BiSeNet(n_classes=19).to(self.device)
        if use_cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count()))

        print("\nload the weight from pretrained-weight file")
        model_dict = self.model.state_dict()
        checkpoint = torch.load(self.args.pretrained_weight)['model_state_dict']
        model_dict.update(checkpoint)
        self.model.load_state_dict(model_dict, strict=True)
        print("Restoring the weight from pretrained-weight file \nFinished loading the weight\n")
        self.model.eval()

    @torch.no_grad()
    def predict(self, image_path):
        image = Image.open(image_path).convert("RGB")
        width = image.width
        height = image.height

        data = predict_transform(image)
        data = data.unsqueeze(0).to(self.device)
        print(data.shape)
        output, output16, output32 = self.model(data)

        pred = torch.argmax(output16, dim=1, keepdim=False).squeeze(0).cpu().data.numpy()

        result = np.zeros([height, width, 3])
        for i in range(len(part_colors)):
            index = np.where(pred == i)
            result[index[0], index[1], :] = part_colors[i]

        mask = Image.fromarray(result.astype(np.uint8))
        new_image = Image.blend(image, mask, 0.5)
        new_image.show()
        new_image.save("result.png")


if __name__ == "__main__":
    model = Inference()

    # model.predict("E:/Datasets2/CelebAMask-HQ/CelebA-HQ-img/1.jpg")
    model.predict("1.jpg")